import matplotlib.pyplot as plt
import numpy as np
import math
def draw(dots, connects):
    plt.figure(figsize=(6,6))
    plt.xlim(0,101)     
    plt.ylim(0,101)       
    plt.xlabel('X')    
    plt.ylabel('Y')    
    for i in range(len(dots)):
        plt.text(dots[i][0]+ 0.5,dots[i][1]+0.5,str(i),color='#0085c3')
        plt.plot(dots[i][0], dots[i][1],'o',color='#ff1493')
    for i in connects:
        start = (dots[i[0]][0], dots[i[1]][0])
        end = (dots[i[0]][1], dots[i[1]][1])
        plt.plot(start, end, color='#ff1493')
    plt.title("Prim Algorithm - Minimum Spining Tree")
    plt.grid()
    plt.show()

def generate_graph(Points_Num):
    NaivePoints = np.random.rand(Points_Num,2)
    points = []
    for i in NaivePoints:
        points.append([i[0] * Points_Num * 2, i[1] * Points_Num * 2])
    graph = dict()
    size = len(points)
    for i in range(size):
        x1, y1 = points[i][0], points[i][1]
        for j in range(size):
            x2, y2 =  points[j][0], points[j][1]
            dist = math.sqrt((x2 - x1)**2 + (y2 - y1)**2)
            if i not in graph:
                graph[i] = dict()
            if j not in graph:
                graph[j] = dict()
            graph[i][j] = dist
            graph[j][i] = dist
    return points, graph

def prim_algorithm(graph):
    size = len(graph)
    vis = set()
    dist = [float('inf') for _ in range(size)]
    ans = 0         # 最小生成树长度
    pos = 0
    dist[pos] = 0   # 最短邻边长度
    vis.add(pos)    # 是否访问
    result = []     # 存储结果
    pre = [float('inf') for _ in range(size)]

    for i in range(1, size):
        if 0 in graph and i in graph[0]:
            dist[i] = graph[0][i]
            pre[i] = 0
    for i in range(size - 1):
        cur_min = float('inf')
        pos = -1
        for j in range(size):
            if j not in vis and dist[j] < cur_min:
                cur_min = dist[j]
                pos = j
        if pos == -1:
            return -1
        ans += cur_min
        result.append([pos, pre[pos]])
        vis.add(pos)
        for j in range(size):
            if j not in vis and dist[j] > graph[pos][j]:
                dist[j] = graph[pos][j]
                pre[j] = pos
    return ans, result

if __name__ == "__main__":
    points , Graph = generate_graph(50)
    ans , result = prim_algorithm(Graph)
    draw(points, result)

    
